# 导入json模块，用于处理JSON数据
import json

# 导入logging模块，用于日志记录
import logging
import threading

# 导入requests模块，用于发送HTTP请求
import requests

# 从urllib.parse模块导入urljoin，用于拼接URL
from urllib.parse import urljoin

# 配置日志级别为INFO
logging.basicConfig(level=logging.INFO)
# 获取当前模块的日志记录器
logger = logging.getLogger(__name__)


# 定义SSE监听器函数，负责持续监听服务器推送的事件
def start_sse_listerner(client):
    # 打开SSE流并持续监听
    sse_response = client.open_sse_stream()
    if sse_response:
        print("✅ SSE流连接成功")
        try:
            # 持续监听SSE事件
            for line in sse_response.iter_lines(chunk_size=1, decode_unicode=True):
                # 处理event事件类型行
                if line.startswith("event: "):
                    event_type = line[7:]  # 去掉 'event: ' 前缀
                # 处理data数据行
                elif line.startswith("data: "):
                    data_str = line[6:]  # 去掉 'data: ' 前缀
                    try:
                        # 解析JSON数据
                        data = json.loads(data_str)
                        # 处理connected事件
                        if event_type == "connected":
                            status = data.get("status")
                            session_id = data.get("sessionId")
                            print(f"🔗 连接状态: {status}, 会话ID: {session_id}")
                        # 处理time_update事件
                        elif event_type == "time_update":
                            server_time = data.get("server_time")
                            timestamp = data.get("timestamp")
                            session_id = data.get("sessionId")
                            print(
                                f"🕐 服务器时间: {server_time} (时间戳: {timestamp}, 会话ID: {session_id})"
                            )
                        # 处理未知事件类型
                        else:
                            print(f"📡 未知事件类型: {event_type}, 数据: {data}")
                    # 捕获JSON解析异常
                    except json.JSONDecodeError:
                        print(f"❌ 数据解析错误: {data_str}")
                # 空行表示事件结束，跳过
                elif line == "":
                    continue
        # 捕获用户中断
        except KeyboardInterrupt:
            print("\n⏹️  用户停止监听")
        finally:
            # 关闭SSE流
            sse_response.close()
            print("🔌 SSE流已关闭")


# 定义MCPClient类，表示MCP HTTP客户端
class MCPClient:
    """MCP HTTP客户端"""

    # 初始化方法，设置服务器URL和相关属性
    def __init__(self, server_url: str):
        # 协议版本号
        self.protocol_version = "2025-06-18"
        # 客户端信息
        self.client_info = {"name": "MCP客户端", "version": "1.0.0"}
        # 服务器信息，初始化为None
        self.server_info = None
        # 服务器能力，初始化为None
        self.capabilities = None
        # 会话ID，初始化为None
        self.session_id = None
        # 消息ID，初始值为1
        self.message_id = 1

        # 去除服务器URL末尾的斜杠
        self.server_url = server_url.rstrip("/")
        # 拼接MCP接口的完整URL
        self.mcp_endpoint = urljoin(self.server_url, "/mcp")
        # 创建requests的Session对象
        self.session = requests.Session()
        # 设置Session的默认请求头
        self.session.headers.update(
            {
                "Accept": "application/json, text/event-stream",
                "MCP-Protocol-Version": self.protocol_version,
            }
        )

    # 重置会话方法
    def reset_session(self):
        """重置会话"""
        # 保存旧的会话ID
        old_session_id = self.session_id
        # 清空当前会话ID
        self.session_id = None
        # 记录会话重置信息
        logger.info(f"会话已重置，原会话ID: {old_session_id}")

    # 打开SSE流连接方法
    def open_sse_stream(self):
        """打开SSE流连接"""
        try:
            # 构造请求头
            headers = {"Accept": "text/event-stream"}
            # 如果有会话ID则加入请求头
            if self.session_id:
                headers["Mcp-Session-Id"] = self.session_id

            # 发送GET请求，开启SSE流
            response = self.session.get(
                self.mcp_endpoint, headers=headers, stream=True, timeout=30
            )

            # 如果响应状态码为200，表示连接成功
            if response.status_code == 200:
                logger.info("SSE流已打开")
                return response
            else:
                # 连接失败，记录错误
                logger.error(f"SSE流打开失败: {response.status_code}")
                return None

        # 捕获请求异常
        except requests.exceptions.RequestException as e:
            logger.error(f"SSE流连接错误: {e}")
            return None

    # 获取当前会话信息方法
    def get_session_info(self):
        """获取当前会话信息"""
        # 返回会话相关信息字典
        return {
            "session_id": self.session_id,
            "server_url": self.server_url,
            "protocol_version": self.protocol_version,
        }

    # 获取可用工具列表方法
    def list_tools(self):
        """获取可用工具列表"""
        # 构造请求消息
        request = {
            "jsonrpc": "2.0",
            "id": self._generate_id(),
            "method": "tools/list",
            "params": {},
        }

        # 发送请求并获取响应
        response = self._send_message(request)
        # 如果响应存在且包含result字段
        if response and "result" in response:
            # 获取工具列表
            tools = response["result"].get("tools", [])
            logger.info(f"可用工具数量: {len(tools)}")
            return tools
        else:
            # 获取失败，记录错误
            logger.error("获取工具列表失败")
            return []

    # 通过HTTP发送消息的私有方法
    def _send_message(self, message):
        """通过HTTP发送消息"""
        try:
            # 初始化请求头字典
            headers = {}
            # 如果有会话ID，则添加到请求头
            if self.session_id:
                headers["Mcp-Session-Id"] = self.session_id
                logger.debug(f"使用会话ID: {self.session_id}")

            # 发送POST请求
            response = self.session.post(
                self.mcp_endpoint, json=message, headers=headers, timeout=30
            )

            # 检查响应状态码
            if response.status_code == 200:
                # 检查响应头中是否有新的会话ID
                new_session_id = response.headers.get("Mcp-Session-Id")
                # 如果有新的会话ID且与当前不同，则更新
                if new_session_id and new_session_id != self.session_id:
                    self.session_id = new_session_id
                    logger.info(f"更新会话ID: {self.session_id}")

                # 返回响应的JSON内容
                return response.json()
            # 如果状态码为202，表示已接受但无内容
            elif response.status_code == 202:
                # Accepted - 通知或响应
                return None
            # 如果状态码为400，表示请求错误
            elif response.status_code == 400:
                logger.error(f"请求错误: {response.text}")
                return None
            # 如果状态码为401，表示会话ID无效
            elif response.status_code == 401:
                logger.error("会话ID无效，需要重新初始化")
                self.session_id = None
                return None
            # 其他HTTP错误
            else:
                logger.error(f"HTTP错误: {response.status_code} - {response.text}")
                return None

        # 捕获请求异常
        except requests.exceptions.RequestException as e:
            logger.error(f"HTTP请求错误: {e}")
            return None

    # 生成消息ID的方法
    def _generate_id(self) -> str:
        """生成消息ID"""
        # 消息ID自增
        self.message_id += 1
        # 返回字符串类型的消息ID
        return str(self.message_id)

    # 发送消息到服务器的方法（重复定义，实际只会用后面这个）
    def _send_message(self, message):
        """通过HTTP发送消息"""
        try:
            # 初始化请求头字典
            headers = {}
            # 如果有会话ID，则添加到请求头
            if self.session_id:
                headers["Mcp-Session-Id"] = self.session_id
                logger.debug(f"使用会话ID: {self.session_id}")

            # 发送POST请求到MCP端点
            response = self.session.post(
                self.mcp_endpoint, json=message, headers=headers, timeout=30
            )

            # 检查响应状态码
            if response.status_code == 200:
                # 从响应头获取新的会话ID
                new_session_id = response.headers.get("Mcp-Session-Id")
                # 如果有新的会话ID且与当前不同，则更新
                if new_session_id and new_session_id != self.session_id:
                    self.session_id = new_session_id
                    logger.info(f"更新会话ID: {self.session_id}")

                # 返回响应的JSON内容
                return response.json()
            # 如果状态码为202，表示已接受但无内容
            elif response.status_code == 202:
                # Accepted - 通知或响应
                return None
            # 如果状态码为400，表示请求错误
            elif response.status_code == 400:
                logger.error(f"请求错误: {response.text}")
                return None
            # 如果状态码为401，表示会话ID无效
            elif response.status_code == 401:
                logger.error("会话ID无效，需要重新初始化")
                self.session_id = None
                return None
            # 其他HTTP错误
            else:
                logger.error(f"HTTP错误: {response.status_code} - {response.text}")
                return None

        # 捕获请求异常
        except requests.exceptions.RequestException as e:
            logger.error(f"HTTP请求错误: {e}")
            return None

    # 调用工具方法
    def call_tool(self, tool_name, arguments):
        """调用工具"""
        # 构造调用工具的请求消息
        request = {
            "jsonrpc": "2.0",
            "id": self._generate_id(),
            "method": "tools/call",
            "params": {"calls": [{"name": tool_name, "arguments": arguments}]},
        }

        # 发送请求并获取响应
        response = self._send_message(request)
        # 如果响应存在且包含result字段
        if response and "result" in response:
            # 获取调用结果
            calls = response["result"].get("calls", [])
            if calls:
                return calls[0]
        else:
            # 调用失败，记录错误
            logger.error(f"调用工具 {tool_name} 失败")
        return None

    # 初始化与服务器连接的方法
    def initialize(self) -> bool:
        """初始化与服务器的连接"""
        # 构造初始化请求消息
        request = {
            "jsonrpc": "2.0",
            "id": self._generate_id(),
            "method": "initialize",
            "params": {
                "protocolVersion": self.protocol_version,
                "capabilities": {},
                "clientInfo": self.client_info,
            },
        }

        # 发送初始化请求
        response = self._send_message(request)
        # 如果响应存在且包含result字段
        if response and "result" in response:
            result = response["result"]
            # 获取服务器信息
            self.server_info = result.get("serverInfo")
            # 获取服务器能力
            self.capabilities = result.get("capabilities")

            # 提取并存储会话ID
            if "sessionId" in result:
                self.session_id = result["sessionId"]
                logger.info(f"获取到会话ID: {self.session_id}")

            logger.info(f"服务器初始化成功: {self.server_info}")
            return True
        else:
            # 初始化失败，记录错误
            logger.error("服务器初始化失败")
            return False


# HTTP传输的函数
def http_transport():
    """HTTP传输"""
    # 打印分隔线
    print("=== HTTP传输 ===")

    # 创建HTTP客户端实例
    client = MCPClient("http://localhost:8000")

    try:
        # 初始化客户端与服务器的连接
        if not client.initialize():
            print("初始化失败")
            return
        # 启动SSE监听线程
        sse_listerner_thread = threading.Thread(
            target=start_sse_listerner, args=(client,), daemon=True
        )
        sse_listerner_thread.start()
        # 显示会话信息
        session_info = client.get_session_info()
        print(f"✅ 会话信息: {session_info}")
        # 获取工具列表
        tools = client.list_tools()
        print(f"可用工具: {[tool['name'] for tool in tools]}")

        # 调用calculate工具
        result = client.call_tool("calculate", {"expression": "10 / 2 + 5"})
        if result:
            # 获取内容字段
            content = result.get("content", [])
            if content:
                print(f"计算结果: {content[0].get('text', '')}")

        # 最终会话状态
        final_session_info = client.get_session_info()
        print(f"🏁 最终会话状态: {final_session_info}")

        # 会话重置
        print("\n=== 会话重置 ===")
        client.reset_session()
        reset_session_info = client.get_session_info()
        print(f"🔄 重置后会话状态: {reset_session_info}")
        # 等待SSE监听线程结束
        sse_listerner_thread.join()
    # 捕获并打印异常
    except Exception as e:
        print(f"HTTP错误: {e}")


# 主函数
def main():
    """主函数"""
    # 调用HTTP传输函数
    http_transport()


# 判断是否为主程序入口
if __name__ == "__main__":
    main()
